package com.qyer.commons.kafka;

import com.qyer.commons.utils.ReentrantExceptionOperator;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
import org.apache.kafka.common.TopicPartition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Arrays;
import java.util.Collection;
import java.util.concurrent.CountDownLatch;

/**
 * User: Z J Wu Date: 2016/06/28 Time: 18:10 Package: com.qyer.pushengine.ios.model
 */
public class CommitWhenRebalanceListener implements ConsumerRebalanceListener {

  private static final Logger LOGGER = LoggerFactory.getLogger(CommitWhenRebalanceListener.class);

  public CommitWhenRebalanceListener(ReentrantKafkaCommit rc,
                                     ReentrantExceptionOperator<Void> operator,
                                     PartitionAssignment assignment, CountDownLatch initSignal) {
    this.rc = rc;
    this.operator = operator;
    this.assignment = assignment;
    this.initSignal = initSignal;
  }

  private final ReentrantKafkaCommit rc;
  private final ReentrantExceptionOperator<Void> operator;

  private final PartitionAssignment assignment;

  private boolean firstTimeRebalance = true;

  private CountDownLatch initSignal;

  @Override
  public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
    LOGGER.info("Consumers changed, will rebalance now.");
    rc.prepareCommit();
    try {
      operator.tryOperate();
      LOGGER.info("Commit for rebalance finished.");
    } catch (Exception e) {
      LOGGER.info("Error when recommit before rebalance.", e);
    } finally {
      operator.clear();
    }
  }

  @Override
  public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
    LOGGER.info("Consumers changed, prepare to rebalance.");
    try {
      if (CollectionUtils.isEmpty(partitions)) {
        LOGGER.warn("Kafka does not assign any partition to current consumer(Empty array).");
        return;
      }
      TopicPartition[] as = partitions.toArray(new TopicPartition[partitions.size()]);
      assignment.setAssignment(as);
      LOGGER
        .info("Rebalance finished. Current partitions for consumer is ({})", Arrays.toString(as));
    } finally {
      if (firstTimeRebalance) {
        initSignal.countDown();
        firstTimeRebalance = false;
      }
    }

  }
}
